##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(RColorBrewer)
library(ggpubr)
library(ggsci)
library(patchwork)

##########################################################################################

option_list <- list(
    make_option(c("--input_path"), type = "character") ,
    make_option(c("--type"), type = "character") ,
    make_option(c("--base_type"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    input_path <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/smgs/mode_baseline_divide"
    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/smgs/mode_baseline_divide_plot"
    type <- "GC"
    base_type <- "Drink"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

input_path <- opt$input_path
out_path <- opt$out_path
type <- opt$type
base_type <- opt$base_type

###########################################################################################

dir.create(out_path , recursive = T)
col <- c( "#006699","#DDA520"  )

###########################################################################################
## 输入文件
input_files <- grep( type , grep( "csv" , grep(base_type , list.files(input_path) , value = T) , value = T) , value = T )

## 合并
tmp_res <- c()
for( file in input_files ){
    tmp <- data.frame(fread(paste0( input_path , "/" , file)))
    tmp_res <- rbind(tmp_res , tmp)
}
if(type=="GC"){
   tmp_res$gene_type <- ifelse(tmp_res$gene == "CDH1" , "Maintained" , tmp_res$gene_type )
}
tmp_res <- tmp_res %>%
group_by( gene , variable , gene_type , base_type ) %>%
summarize( value = sum(value) )

## 计算两组是否分布存在差异
trans <- function(num){
    up <- floor(log10(num))
    down <- round(num*10^(-up),2)
    text <- paste("P == ",down," %*% 10","^",up)
    return(text)
}

base_count <- subset(tmp_res , gene == "SMG") %>%
group_by(base_type) %>%
summarize(count_all = sum(value))
tmp_res <- merge(tmp_res , base_count , by = "base_type")
tmp_res$ratio <- tmp_res$value/tmp_res$count_all
tmp_res$value_text <- round(tmp_res$ratio , 2) * 100
tmp_res <- tmp_res[order(tmp_res$variable),]

result <- c()
tmp_res <- subset(tmp_res , gene == "SMG")
for( geneN in unique(tmp_res$gene) ){

    tmp <- subset( tmp_res , gene == geneN )
    ## 分两组进行检验
    tmp1 <- subset( tmp , base_type == unique(tmp$base_type)[1] )
    tmp2 <- subset( tmp , base_type == unique(tmp$base_type)[2] )
    ## 在两组存在突变
    if( sum(tmp1$value_text)!=0 || sum(tmp2$value_text)!=0 ){
        p <- chisq.test(matrix(c(tmp1$value , tmp2$value) , ncol = 2))$p.value
        if( p < 0.01 ){
            p_text <- trans(p)
        }else{
            p_text <- paste0( "P == " , round(as.numeric(p) , 2) ) 
        }
        tmp$p_text <- ""
        tmp$p_text[1] <- p_text
        result <- rbind(result , tmp)
    }
}
#print(result)

###########################################################################################
## 基因分布堆叠图
if(1!=1){
    dat <- result
    dat$variable <- factor( dat$variable , levels = c("IM branch" , "GC branch" , "Share") , order = T )
    dat$gene_type <- factor( dat$gene_type , levels = c("Maintained" , "IM_favored" , "GC_favored") , order = T )
    gene_order <- c("TP53","ARID1A","CDH1","APC","SMAD4","MUC6","PIK3CA",
                  "CTNNB1","RHOA","ERBB2","CFTR","KRAS","MAP2K7","ARID2",
                  "RNF43","TGFBR2","BMP6","FBXW7","CDKN2A","MTRR")
    dat$gene <- factor( dat$gene , levels = gene_order , order = T )
    dat$base_type <- sapply(strsplit(dat$base_type,"_"),"[" , 2)
    dat$base_type <- factor( dat$base_type , levels = c("Drinker" , "Nondrinker" , "Postive" , "Negative" , "smoker" , "nonsmoker" , "Younger" , "Older") )

    col_use <- c(rgb(red=179,green=34,blue=35,alpha=255,max=255) ,
        rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
        rgb(red=2,green=100,blue=190,alpha=255,max=255) ,
        "#4DAF4A"
        )

    col_use <- col_use[c(4,3,1)]
    names(col_use) <- c("IM branch" , "GC branch" , "Share")

    ## 构成比
    plotFunc <- function(dat = dat , gene_typeN = gene_typeN){
        tmp <- dat %>% group_by(gene) %>% summarize(mutnum=sum(value))
        y_max <- max(tmp$mutnum) * 1.1

        p1 <- ggplot(dat , aes(base_type, weight= value , fill=variable))+
        geom_bar(position="stack")+
       # title(gene_typeN) +
        labs(y="Number of patients")+
        facet_grid(.~gene , scales = "free" ) +
        geom_text(aes(label=p_text , y = y_max , x = 1.5),parse = TRUE,size=3 , color = "black", face='bold') +
        theme(panel.background = element_rect(fill = NA, colour = "black", size = 1),
            legend.position ='top',
            panel.grid.major = element_line(colour=NA),
            legend.text = element_text(size = 8,color="black",face='bold'),
            axis.text.y = element_text(size = 8,color="black",face='bold'),
            axis.title.y = element_text(size = 12,color="black",face='bold'),
            axis.title.x=element_blank(),
            #axis.text.x = element_text(size = 12,color="black",face='bold' , angle = 90 ,  vjust = .5, hjust = .5),
            axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1 , size = 12,color="black",face='bold'),
            strip.text.x = element_text(size = 8 , face = 'bold'),
            axis.ticks.length = unit(0.2, "cm") ,
            axis.line = element_line(size = 0.5)) +
        guides(fill = guide_legend(reverse=TRUE , title = NULL)) +
        scale_fill_manual(values = col_use)
        p <- p1
        gp <- ggplotGrob(p)
        facet.columns <- gp$layout$l[grepl("panel", gp$layout$name)]
        x.var <- sapply(ggplot_build(p)$layout$panel_scales_x,
                        function(l) length(l$range$range))
        gp$widths[facet.columns] <- gp$widths[facet.columns] * x.var
        
        width = length(unique(dat$gene))/1.4 + 1
        height= 4 
        image_name <- paste0(out_path , "/GeneMode.",type,".",base_type,".",gene_typeN,".pdf")
        pdf(image_name , width=width ,height=height)
        grid::grid.draw(gp)
        dev.off()
    }

    gene_typeN <- "Maintained"
    plotFunc(dat = subset(dat , gene_type == gene_typeN) , gene_type = gene_typeN)
    gene_typeN <- "IM_favored"
    plotFunc(dat = subset(dat , gene_type == gene_typeN) , gene_type = gene_typeN)
    gene_typeN <- "GC_favored"
    plotFunc(dat = subset(dat , gene_type == gene_typeN) , gene_type = gene_typeN)
}

###########################################################################################
## 基因分布堆叠图
# 只画SMG的
dat <- subset(result , gene=="SMG")
#base_count <- dat %>%
#group_by(base_type) %>%
#summarize(count_all = sum(value))
#dat <- merge(dat , base_count , by = "base_type")
#dat$ratio <- dat$value/dat$count_all

dat$variable <- factor( dat$variable , levels = c("IM branch" , "GC branch" , "Share") , order = T )
dat$base_type <- sapply(strsplit(dat$base_type,"_"),"[" , 2)
dat$base_type <- factor( dat$base_type , levels = c("Drinker" , "Nondrinker" , "Postive" , "Negative" , "smoker" , "nonsmoker" , "Younger" , "Older") )
dat$ratio[is.na(dat$ratio)] <- 0
dat$value_text <- round(dat$ratio , 2) * 100
dat$value_text <- ifelse( dat$ratio == 0 , "" , dat$value_text )


col_use <- c(rgb(red=179,green=34,blue=35,alpha=255,max=255) ,
    rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
    rgb(red=2,green=100,blue=190,alpha=255,max=255) ,
    "#4DAF4A"
    )

col_use <- col_use[c(4,3,1)]
names(col_use) <- c("IM branch" , "GC branch" , "Share")

## 构成比
plot <- ggplot( data = dat , aes( x = base_type , y = ratio , fill = variable ))+
    geom_bar(position = "stack", stat = "identity") + 
    theme_bw()+
    labs(x="",y="Proportion (%)")+
    theme(panel.grid = element_blank())+
    scale_fill_npg() +
    ylim(0,1.05)+
    geom_text(aes(label=p_text , y = 1.05 ,x = 1.5),parse = TRUE,size=4)+
    geom_text(aes(label=value_text) , position=position_stack(vjust = 0.5) , size=3 , color="black")+
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
                legend.position ='right',
                legend.title = element_blank() ,
                panel.grid.major=element_line(colour=NA),
                plot.title = element_text(size = 12,color="black",face='bold'),
                legend.text = element_text(size = 8,color="black",face='bold'),
                axis.text.y = element_text(size = 7,color="black",face='bold'),
                axis.title.x = element_text(size = 12,color="black",face='bold'),
                axis.title.y = element_text(size = 12,color="black",face='bold'),
                strip.text.x = element_text(size = 7 , face = 'bold'),
                axis.text.x = element_text(size = 8,color="black",face='bold') ,
                axis.line = element_line(size = 0.5))

image_name <- paste0(out_path , "/GeneMode.",type,".",base_type,".SMG.pdf")
ggsave( image_name , plot , width = 3 , height = 3 )
